import numpy as np
import pandas as pd
from scipy.interpolate import interp1d

def read_truth_data(filename):
    data = np.loadtxt(filename)
    wavelengths = data[:, 0]
    intensities = data[:, 1]
    return wavelengths, intensities

def read_reconstructed_csv(filename):
    df = pd.read_csv(filename)
    wavelengths = df.iloc[:, 0].values
    intensities = df.iloc[:, 1].values
    return wavelengths, intensities

def calculate_psnr(truth_wl, truth_intensity, recon_wl, recon_intensity,
                   intensity_threshold_ratio=0.05):
    """
    计算重构光谱相对于真值光谱的PSNR值。
    - 对重构光谱做插值，使其波长与真值对应
    - 只考虑真值强度超过阈值的波长点，避免噪声影响

    返回PSNR值（单位：dB）
    """

    # 插值重构光谱
    interp_func = interp1d(recon_wl, recon_intensity, kind='linear', fill_value='extrapolate')
    recon_interp = interp_func(truth_wl)

    # 选取真值强度较大点
    max_truth_intensity = np.max(truth_intensity)
    valid_idx = truth_intensity >= (max_truth_intensity * intensity_threshold_ratio)
    if np.sum(valid_idx) == 0:
        print("警告：无有效信号点满足强度阈值，无法计算PSNR。")
        return 0

    truth_valid = truth_intensity[valid_idx]
    recon_valid = recon_interp[valid_idx]

    # 计算均方误差（MSE）
    mse = np.mean((truth_valid - recon_valid) ** 2)
    if mse == 0:
        return float('inf')  # 完美重构，PSNR无穷大

    # 计算峰值信号强度，使用真值光谱最大值
    peak_signal = max_truth_intensity

    psnr = 10 * np.log10(peak_signal ** 2 / mse)
    return psnr

def main():
    truth_file = 'truth_spectrum.txt'
    recon_file = 'reconstructed.csv'

    truth_wl, truth_intensity = read_truth_data(truth_file)
    recon_wl, recon_intensity = read_reconstructed_csv(recon_file)

    psnr = calculate_psnr(truth_wl, truth_intensity, recon_wl, recon_intensity,
                          intensity_threshold_ratio=0.05)
    print(f'重构光谱与真值光谱的PSNR为: {psnr:.2f} dB')

    # 根据经验阈值判断重构质量，这里阈值可根据实际需求调整
    if psnr >= 30:
        print("重构成功：PSNR达到30 dB以上")
    else:
        print("重构失败：PSNR不足30 dB")

if __name__ == '__main__':
    main()
